#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
代理池管理模块

管理代理IP池，提供自动切换和健康检测功能。
"""

import asyncio
import aiohttp
import logging
import time
from datetime import datetime, timedelta
from typing import Dict, List, Optional, Tuple
from dataclasses import dataclass
from enum import Enum
import random

class ProxyStatus(Enum):
    """代理状态"""
    ACTIVE = "active"
    INACTIVE = "inactive"
    BANNED = "banned"
    TIMEOUT = "timeout"
    ERROR = "error"

@dataclass
class ProxyInfo:
    """代理信息"""
    host: str
    port: int
    username: str = ""
    password: str = ""
    protocol: str = "http"  # http, https, socks5
    status: ProxyStatus = ProxyStatus.ACTIVE
    success_count: int = 0
    failure_count: int = 0
    last_used: str = ""
    last_check: str = ""
    response_time: float = 0.0
    country: str = ""
    city: str = ""
    provider: str = ""
    
    @property
    def proxy_url(self) -> str:
        """获取代理URL"""
        if self.username and self.password:
            return f"{self.protocol}://{self.username}:{self.password}@{self.host}:{self.port}"
        else:
            return f"{self.protocol}://{self.host}:{self.port}"
    
    @property
    def success_rate(self) -> float:
        """成功率"""
        total = self.success_count + self.failure_count
        return self.success_count / max(total, 1)

class ProxyPool:
    """代理池管理器"""
    
    def __init__(self, config: Dict = None):
        self.logger = logging.getLogger(__name__)
        
        # 默认配置
        self.config = {
            "max_retries": 3,
            "timeout": 10,
            "check_interval": 300,  # 5分钟检查一次
            "min_success_rate": 0.7,
            "max_failure_count": 10,
            "health_check_url": "http://httpbin.org/ip",
            "enable_auto_refresh": True,
            "proxy_rotation": True
        }
        
        if config:
            self.config.update(config)
        
        # 代理池
        self.proxies: List[ProxyInfo] = []
        self.current_proxy_index = 0
        
        # 统计信息
        self.stats = {
            "total_proxies": 0,
            "active_proxies": 0,
            "banned_proxies": 0,
            "total_requests": 0,
            "successful_requests": 0,
            "failed_requests": 0,
            "avg_response_time": 0.0
        }
        
        # 健康检查任务
        self.health_check_task = None
        self.is_running = False
        
        self.logger.info("代理池管理器初始化完成")
    
    def add_proxy(self, host: str, port: int, username: str = "", 
                  password: str = "", protocol: str = "http", **kwargs) -> bool:
        """添加代理"""
        try:
            proxy = ProxyInfo(
                host=host,
                port=port,
                username=username,
                password=password,
                protocol=protocol,
                **kwargs
            )
            
            # 检查是否已存在
            for existing_proxy in self.proxies:
                if (existing_proxy.host == host and 
                    existing_proxy.port == port and
                    existing_proxy.username == username):
                    self.logger.warning(f"代理已存在: {host}:{port}")
                    return False
            
            self.proxies.append(proxy)
            self.stats["total_proxies"] += 1
            
            self.logger.info(f"已添加代理: {host}:{port}")
            return True
            
        except Exception as e:
            self.logger.error(f"添加代理失败: {e}")
            return False
    
    def add_proxy_list(self, proxy_list: List[Dict]) -> int:
        """批量添加代理"""
        added_count = 0
        
        for proxy_data in proxy_list:
            if self.add_proxy(**proxy_data):
                added_count += 1
        
        self.logger.info(f"批量添加代理完成，成功: {added_count}/{len(proxy_list)}")
        return added_count
    
    def remove_proxy(self, host: str, port: int) -> bool:
        """移除代理"""
        for i, proxy in enumerate(self.proxies):
            if proxy.host == host and proxy.port == port:
                removed_proxy = self.proxies.pop(i)
                self.stats["total_proxies"] -= 1
                
                # 调整当前代理索引
                if i <= self.current_proxy_index and self.current_proxy_index > 0:
                    self.current_proxy_index -= 1
                
                self.logger.info(f"已移除代理: {host}:{port}")
                return True
        
        return False
    
    async def get_proxy(self) -> Optional[ProxyInfo]:
        """获取可用代理"""
        if not self.proxies:
            return None
        
        # 筛选可用代理
        available_proxies = [
            proxy for proxy in self.proxies 
            if proxy.status == ProxyStatus.ACTIVE and 
               proxy.success_rate >= self.config["min_success_rate"] and
               proxy.failure_count < self.config["max_failure_count"]
        ]
        
        if not available_proxies:
            # 如果没有可用代理，尝试使用所有代理
            available_proxies = [
                proxy for proxy in self.proxies 
                if proxy.status != ProxyStatus.BANNED
            ]
        
        if not available_proxies:
            self.logger.warning("没有可用的代理")
            return None
        
        # 代理选择策略
        if self.config["proxy_rotation"]:
            # 轮询选择
            proxy = available_proxies[self.current_proxy_index % len(available_proxies)]
            self.current_proxy_index = (self.current_proxy_index + 1) % len(available_proxies)
        else:
            # 随机选择
            proxy = random.choice(available_proxies)
        
        # 更新使用时间
        proxy.last_used = datetime.now().isoformat()
        
        return proxy
    
    async def test_proxy(self, proxy: ProxyInfo) -> Tuple[bool, float]:
        """测试代理可用性"""
        start_time = time.time()
        
        try:
            proxy_dict = {
                "http": proxy.proxy_url,
                "https": proxy.proxy_url
            }
            
            timeout = aiohttp.ClientTimeout(total=self.config["timeout"])
            
            async with aiohttp.ClientSession(timeout=timeout) as session:
                async with session.get(
                    self.config["health_check_url"],
                    proxy=proxy.proxy_url
                ) as response:
                    response_time = time.time() - start_time
                    
                    if response.status == 200:
                        # 验证返回的IP是否为代理IP
                        response_text = await response.text()
                        
                        # 更新代理信息
                        proxy.response_time = response_time
                        proxy.last_check = datetime.now().isoformat()
                        proxy.success_count += 1
                        proxy.status = ProxyStatus.ACTIVE
                        
                        self.logger.debug(f"代理测试成功: {proxy.host}:{proxy.port}, 响应时间: {response_time:.2f}s")
                        return True, response_time
                    else:
                        raise aiohttp.ClientError(f"HTTP {response.status}")
        
        except asyncio.TimeoutError:
            proxy.status = ProxyStatus.TIMEOUT
            proxy.failure_count += 1
            response_time = time.time() - start_time
            self.logger.warning(f"代理超时: {proxy.host}:{proxy.port}")
            return False, response_time
            
        except Exception as e:
            proxy.status = ProxyStatus.ERROR
            proxy.failure_count += 1
            response_time = time.time() - start_time
            
            # 检查是否被封禁
            if "403" in str(e) or "banned" in str(e).lower():
                proxy.status = ProxyStatus.BANNED
                self.logger.warning(f"代理可能被封禁: {proxy.host}:{proxy.port}")
            
            self.logger.error(f"代理测试失败: {proxy.host}:{proxy.port}, 错误: {e}")
            return False, response_time
    
    async def health_check(self):
        """健康检查所有代理"""
        if not self.proxies:
            return
        
        self.logger.info("开始代理健康检查...")
        
        # 并发测试代理
        semaphore = asyncio.Semaphore(10)  # 限制并发数
        
        async def test_single_proxy(proxy: ProxyInfo):
            async with semaphore:
                await self.test_proxy(proxy)
        
        tasks = [test_single_proxy(proxy) for proxy in self.proxies]
        await asyncio.gather(*tasks, return_exceptions=True)
        
        # 更新统计信息
        self.update_stats()
        
        active_count = sum(1 for p in self.proxies if p.status == ProxyStatus.ACTIVE)
        self.logger.info(f"代理健康检查完成，活跃代理: {active_count}/{len(self.proxies)}")
    
    def update_stats(self):
        """更新统计信息"""
        if not self.proxies:
            return
        
        self.stats["total_proxies"] = len(self.proxies)
        self.stats["active_proxies"] = sum(1 for p in self.proxies if p.status == ProxyStatus.ACTIVE)
        self.stats["banned_proxies"] = sum(1 for p in self.proxies if p.status == ProxyStatus.BANNED)
        
        total_requests = sum(p.success_count + p.failure_count for p in self.proxies)
        successful_requests = sum(p.success_count for p in self.proxies)
        
        self.stats["total_requests"] = total_requests
        self.stats["successful_requests"] = successful_requests
        self.stats["failed_requests"] = total_requests - successful_requests
        
        # 平均响应时间
        response_times = [p.response_time for p in self.proxies if p.response_time > 0]
        self.stats["avg_response_time"] = sum(response_times) / len(response_times) if response_times else 0.0
    
    async def start_health_check_task(self):
        """启动健康检查任务"""
        if self.health_check_task and not self.health_check_task.done():
            return
        
        self.is_running = True
        self.health_check_task = asyncio.create_task(self._health_check_loop())
        self.logger.info("代理健康检查任务已启动")
    
    async def _health_check_loop(self):
        """健康检查循环"""
        while self.is_running:
            try:
                await self.health_check()
                await asyncio.sleep(self.config["check_interval"])
            except Exception as e:
                self.logger.error(f"健康检查任务出错: {e}")
                await asyncio.sleep(60)  # 出错后等待1分钟再继续
    
    async def stop_health_check_task(self):
        """停止健康检查任务"""
        self.is_running = False
        
        if self.health_check_task and not self.health_check_task.done():
            self.health_check_task.cancel()
            try:
                await self.health_check_task
            except asyncio.CancelledError:
                pass
        
        self.logger.info("代理健康检查任务已停止")
    
    def get_proxy_stats(self) -> Dict:
        """获取代理统计信息"""
        self.update_stats()
        
        # 添加详细统计
        status_counts = {}
        for status in ProxyStatus:
            status_counts[status.value] = sum(1 for p in self.proxies if p.status == status)
        
        # 响应时间分布
        response_time_ranges = {
            "fast": sum(1 for p in self.proxies if 0 < p.response_time <= 1.0),
            "medium": sum(1 for p in self.proxies if 1.0 < p.response_time <= 3.0),
            "slow": sum(1 for p in self.proxies if p.response_time > 3.0)
        }
        
        # 成功率分布
        success_rate_ranges = {
            "excellent": sum(1 for p in self.proxies if p.success_rate >= 0.9),
            "good": sum(1 for p in self.proxies if 0.7 <= p.success_rate < 0.9),
            "poor": sum(1 for p in self.proxies if p.success_rate < 0.7)
        }
        
        stats = self.stats.copy()
        stats.update({
            "status_distribution": status_counts,
            "response_time_distribution": response_time_ranges,
            "success_rate_distribution": success_rate_ranges,
            "overall_success_rate": self.stats["successful_requests"] / max(self.stats["total_requests"], 1)
        })
        
        return stats
    
    def get_proxy_list(self) -> List[Dict]:
        """获取代理列表"""
        return [
            {
                "host": proxy.host,
                "port": proxy.port,
                "protocol": proxy.protocol,
                "status": proxy.status.value,
                "success_rate": proxy.success_rate,
                "response_time": proxy.response_time,
                "success_count": proxy.success_count,
                "failure_count": proxy.failure_count,
                "last_used": proxy.last_used,
                "last_check": proxy.last_check,
                "country": proxy.country,
                "city": proxy.city,
                "provider": proxy.provider
            }
            for proxy in self.proxies
        ]
    
    def cleanup_failed_proxies(self):
        """清理失败的代理"""
        original_count = len(self.proxies)
        
        # 移除被封禁或失败次数过多的代理
        self.proxies = [
            proxy for proxy in self.proxies
            if not (
                proxy.status == ProxyStatus.BANNED or
                proxy.failure_count >= self.config["max_failure_count"] or
                (proxy.success_rate < self.config["min_success_rate"] and 
                 proxy.success_count + proxy.failure_count > 10)
            )
        ]
        
        removed_count = original_count - len(self.proxies)
        if removed_count > 0:
            self.logger.info(f"已清理 {removed_count} 个失败代理")
            
            # 重置当前代理索引
            if self.current_proxy_index >= len(self.proxies):
                self.current_proxy_index = 0
    
    def reset_proxy_stats(self, proxy_host: str = None, proxy_port: int = None):
        """重置代理统计"""
        if proxy_host and proxy_port:
            # 重置指定代理
            for proxy in self.proxies:
                if proxy.host == proxy_host and proxy.port == proxy_port:
                    proxy.success_count = 0
                    proxy.failure_count = 0
                    proxy.status = ProxyStatus.ACTIVE
                    break
        else:
            # 重置所有代理
            for proxy in self.proxies:
                proxy.success_count = 0
                proxy.failure_count = 0
                proxy.status = ProxyStatus.ACTIVE
        
        self.logger.info("代理统计已重置")
    
    def load_proxy_list_from_file(self, filepath: str) -> int:
        """从文件加载代理列表"""
        try:
            import json
            
            with open(filepath, 'r', encoding='utf-8') as f:
                proxy_data = json.load(f)
            
            if isinstance(proxy_data, list):
                return self.add_proxy_list(proxy_data)
            else:
                self.logger.error("代理文件格式错误，应为列表格式")
                return 0
                
        except Exception as e:
            self.logger.error(f"从文件加载代理失败: {e}")
            return 0
    
    def save_proxy_list_to_file(self, filepath: str) -> bool:
        """保存代理列表到文件"""
        try:
            import json
            
            proxy_data = self.get_proxy_list()
            
            with open(filepath, 'w', encoding='utf-8') as f:
                json.dump(proxy_data, f, ensure_ascii=False, indent=2)
            
            self.logger.info(f"代理列表已保存到: {filepath}")
            return True
            
        except Exception as e:
            self.logger.error(f"保存代理列表失败: {e}")
            return False
    
    async def __aenter__(self):
        """异步上下文管理器入口"""
        if self.config["enable_auto_refresh"]:
            await self.start_health_check_task()
        return self
    
    async def __aexit__(self, exc_type, exc_val, exc_tb):
        """异步上下文管理器出口"""
        await self.stop_health_check_task()
    
    def __len__(self):
        """返回代理池大小"""
        return len(self.proxies)
    
    def __bool__(self):
        """返回是否有代理可用"""
        return len(self.proxies) > 0